#include "core/main.h"

void GameModel::Save(const string &filename) const
{
    OutputDataStream stream;

    stream << _game->Name();
    stream << _parameters;

    const UINT spriteModelCount = _spriteModels.size();
    stream << spriteModelCount;
    for(UINT spriteModelIndex = 0; spriteModelIndex < spriteModelCount; spriteModelIndex++)
    {
        const GameSpriteModel &spriteModel = *_spriteModels[spriteModelIndex];
        stream << spriteModel.spriteAtTree;
        stream << spriteModel.forceStasis;
    }

    const UINT variableModelCount = _variableModels.size();
    stream << variableModelCount;
    for(UINT variableModelIndex = 0; variableModelIndex < variableModelCount; variableModelIndex++)
    {
        const GameVariableModel &variableModel = *_variableModels[variableModelIndex];
        stream << variableModel.trees;
    }

    stream.SaveToFile(filename);
}

void GameModel::Load(Game &game, const string &filename)
{
    InputDataStream stream;
    stream.LoadFromFile(filename);

    _game = &game;
    string gameName;
    stream >> gameName;
    Assert(game.Name() == gameName, "Game name mismatch");

    stream >> _parameters;

    UINT spriteModelCount;
    stream >> spriteModelCount;
    _spriteModels.resize(spriteModelCount);
    for(UINT spriteModelIndex = 0; spriteModelIndex < spriteModelCount; spriteModelIndex++)
    {
        GameSpriteModel *spriteModel = new GameSpriteModel;
        stream >> spriteModel->spriteAtTree;
        stream >> spriteModel->forceStasis;
        _spriteModels[spriteModelIndex] = spriteModel;
    }

    UINT variableModelCount;
    stream >> variableModelCount;
    _variableModels.resize(variableModelCount);
    for(UINT variableModelIndex = 0; variableModelIndex < variableModelCount; variableModelIndex++)
    {
        GameVariableModel *variableModel = new GameVariableModel;
        stream >> variableModel->trees;
        _variableModels[variableModelIndex] = variableModel;
    }

    if(features) delete features;
    features = new short[_parameters.totalVariableCount];
}

struct VariableValue
{
    const GameVariableTree *variable;
    double value;
};

bool operator < (const VariableValue &a, const VariableValue &b)
{
    return a.value > b.value;
}

void GameModel::AdvanceFrame(const vector<const Frame*> &history, Frame &nextFrame) const
{
    const UINT frameCount = history.size();
    vector< unique_ptr<GameModelFrame> > processedFrames(frameCount);

    for(UINT frameIndex = 0; frameIndex < frameCount; frameIndex++)
    {
        processedFrames[frameIndex] = unique_ptr<GameModelFrame>(new GameModelFrame(history[frameIndex]));
        processedFrames[0]->historyFrames.push_back(processedFrames[frameIndex].get());
        processedFrames[frameIndex]->ComputeSpriteCounts(_parameters);
    }

    processedFrames[0]->MakeFeatureVector(_parameters, -9999, -9999, features);

    nextFrame.sprites.clear();
    
    LearnerInput input;
    input.data = features;
    input.length = _parameters.totalVariableCount;

    nextFrame.observedInts.resize(_parameters.observedIntegerCount);
    for(UINT variableIndex = 0; variableIndex < _parameters.observedIntegerCount; variableIndex++)
    {
        const GameVariableModel &curModel = *_variableModels[variableIndex];
        vector<VariableValue> allValues;
        for(UINT treeIndex = 0; treeIndex < curModel.trees.size(); treeIndex++)
        {
            VariableValue curValue;
            curValue.variable = &curModel.trees[treeIndex];
            curValue.value = curValue.variable->tree.EvaluateProbability(input);
            allValues.push_back(curValue);
        }

        std::sort(allValues.begin(), allValues.end());

        const GameVariableTree &chosenVariable = *allValues[0].variable;
        if(chosenVariable.absolute)
        {
            nextFrame.observedInts[variableIndex] = chosenVariable.value;
        }
        else
        {
            nextFrame.observedInts[variableIndex] = history[0]->observedInts[variableIndex] + chosenVariable.value;
        }
    }

    for(int y = 0; y <= _parameters.dimensions.y / _parameters.gridScale; y++)
    {
        for(int x = 0; x <= _parameters.dimensions.x / _parameters.gridScale; x++)
        {
            processedFrames[0]->OverwriteDynamicFeatures(_parameters, x, y, features, false);
            
            for(UINT spriteID = 0; spriteID < _parameters.totalSpriteCount; spriteID++)
            {
                const GameSpriteModel &curModel = *_spriteModels[spriteID];
                double probability = curModel.spriteAtTree.EvaluateProbability(input);
                if(curModel.forceStasis)
                {
                    if(processedFrames[0]->SpriteAtLocation(Vec2i(x, y) * _parameters.gridScale, spriteID) && probability == 0.0)
                    {
                        nextFrame.sprites.push_back(SpriteLocation(Vec2i(x, y) * _parameters.gridScale, spriteID, 0));
                    }
                }
                else
                {
                    if(probability == 0.0)
                    {
                        nextFrame.sprites.push_back(SpriteLocation(Vec2i(x, y) * _parameters.gridScale, spriteID, 0));
                    }
                }
            }
        }
    }
}